﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.Schemes;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;
using Lists = Algoritmia.DataStructures.Lists;

namespace Algoritmia.Problems.SubsetSum
{
    public class SubsetSumStateSpace2 : IReversibleForwardStateSpace<SubsetSumStateSpace2.State, int>
    {
        private readonly Func<IEnumerable<int>, Lists.IList<int>> _createList;
        private readonly Lists.IList<int> _w;
        private readonly int _W;
        private readonly int _n;

        public class State : Lists.ArrayList<int>
        {
            public int Index;
            public int Sum;
            public readonly Lists.IList<int> TailSum;
            public readonly int N;

            public State(Lists.IList<int> w, int N)
                : base(Enumerable.Repeat(0,N))
            {
                Index = 0;
                this.N = N;
                Sum = 0;
                TailSum = Enumerable.Repeat(0, N + 1).ToAlgoritmiaList();
                
                for(int i = N-1; i >= 0; i--)
                {
                    TailSum[i] = TailSum[i + 1] + w[i];
                }

            }
        }

        public SubsetSumStateSpace2(IEnumerable<int> w, int W, Func<IEnumerable<int>, Lists.IList<int>> createList = null)
        {
            _createList = createList ?? (it => AlgoritmiaFactory.Lists<int>.BuildDefault(it));
            _w = _createList(w);
            _W = W;
            _n = _w.Count;
        }

        public State Decide(State state, int decision)
        {
            state[state.Index] = decision;
            state.Sum += decision*_w[state.Index];
            state.Index++;
            return state;
        }

        public IEnumerable<int> Decisions(State state)
        {
            if (state.Index < _n)
            {
                if (state.Sum < _W)
                {
                    if (state.Sum + _w[state.Index] + state.TailSum[state.Index + 1] >= _W)
                    {
                        yield return 1;
                    }
                    if (state.Sum + state.TailSum[state.Index + 1] >= _W)
                    {
                        yield return 0;
                    }
                }
            }
        }

        IEnumerable<State> IForwardStateSpace<State, int>.InitialStates
        {
            get { yield return new State(_w, _n); }
        }

        public bool IsFinal(State state)
        {
            return state.Sum == _W;
        }

        public State Undo(State state, int decision)
        {
            state.Index--;
            state.Sum -= decision*_w[state.Index];
            return state;
        }
    }

    public class SubsetSumSolver2
    {
        private readonly BacktrackingEnumerator<SubsetSumStateSpace2.State, int, Lists.IList<int>> _enumerator;

        public SubsetSumSolver2()
        {
            _enumerator =
                new BacktrackingEnumerator<SubsetSumStateSpace2.State, int, Lists.IList<int>>(
                    (space, initial, decisions, final) =>
                            final.Select((f, i) => i < final.Index ? f : 0).ToAlgoritmiaList()
                        );
        }

        public Lists.IList<int> Solve(IEnumerable<int> w, int W)
        {
            var space = new SubsetSumStateSpace2(w, W);
            return _enumerator.First(space);
        }
    }
}
